# 配置日志
import asyncio
import logging
import signal
import time
from contextlib import asynccontextmanager
from typing import Optional, Dict

from protocol import MessageBuffer, MessageProtocol, MessageType

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(lineno)d - %(message)s'
)
logger = logging.getLogger(__name__)


class ClientConnection:
    """客户端连接管理"""

    def __init__(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter, client_id: str):
        self.reader = reader
        self.writer = writer
        self.client_id = client_id
        self.connected_at = time.time()
        self.last_ping = time.time()
        self.message_buffer = MessageBuffer()
        self.seq_counter = 0

        # 获取客户端地址
        peername = writer.get_extra_info('peername')
        self.address = f"{peername[0]}:{peername[1]}" if peername else "unknown"

    def get_next_seq_id(self) -> int:
        """获取下一个序列号"""
        self.seq_counter += 1
        return self.seq_counter

    async def send_message(self, message: bytes):
        """发送消息到客户端"""
        try:
            self.writer.write(message)
            await self.writer.drain()
        except Exception as e:
            logger.error(f"Failed to send message to {self.client_id}: {e}")
            raise

    async def close(self):
        """关闭连接"""
        try:
            if not self.writer.is_closing():
                self.writer.close()
                await self.writer.wait_closed()
        except Exception as e:
            logger.error(f"Error closing connection {self.client_id}: {e}")


class AsyncTCPServer:
    """异步TCP服务器"""

    def __init__(self, host: str = 'localhost', port: int = 8888):
        self.host = host
        self.port = port
        self.server: Optional[asyncio.Server] = None
        self.clients: Dict[str, ClientConnection] = {}
        self.protocol = MessageProtocol()
        self.running = False
        self.heartbeat_interval = 30  # 心跳间隔(秒)
        self.heartbeat_task: Optional[asyncio.Task] = None

    async def start_server(self):
        """启动服务器"""
        try:
            self.server = await asyncio.start_server(
                self.handle_client,
                self.host,
                self.port
            )
            self.running = True

            # 启动心跳检测任务
            self.heartbeat_task = asyncio.create_task(self.heartbeat_monitor())

            addr = self.server.sockets[0].getsockname()
            logger.info(f"Server started on {addr[0]}:{addr[1]}")

            # 设置信号处理器
            self.setup_signal_handlers()

        except Exception as e:
            logger.error(f"Failed to start server: {e}")
            raise

    async def handle_client(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
        """处理客户端连接"""
        # 生成客户端ID
        peername = writer.get_extra_info('peername')
        client_id = f"client_{peername[0]}_{peername[1]}_{int(time.time())}"

        # 创建客户端连接对象
        client = ClientConnection(reader, writer, client_id)
        self.clients[client_id] = client

        logger.info(f"New client connected: {client_id} from {client.address}")

        try:
            # 发送欢迎消息
            welcome_msg = self.protocol.create_data(
                seq_id=client.get_next_seq_id(),
                data=f"Welcome to PyMicroFlow! Your ID: {client_id}".encode('utf-8')
            )
            await client.send_message(welcome_msg)

            # 处理客户端消息
            while self.running:
                try:
                    # 读取数据 (最多4KB)
                    data = await asyncio.wait_for(reader.read(4096), timeout=60.0)

                    if not data:
                        logger.info(f"Client {client_id} disconnected")
                        break

                    # 解析消息
                    messages = client.message_buffer.feed_data(data)

                    # 处理每个完整消息
                    for message in messages:
                        await self.process_message(client, message)

                except asyncio.TimeoutError:
                    logger.warning(f"Client {client_id} read timeout")
                    break
                except Exception as e:
                    logger.error(f"Error handling client {client_id}: {e}")
                    break

        except asyncio.CancelledError:
            logger.info(f"Client handler {client_id} cancelled")
        except Exception as e:
            logger.error(f"Unexpected error with client {client_id}: {e}")
        finally:
            # 清理客户端连接
            await self.cleanup_client(client_id)

    async def process_message(self, client, message):
        """处理客户端消息"""
        logger.debug(f"Processing message from {client.client_id}: {message.msg_type.name}")

        try:
            if message.msg_type == MessageType.PING:
                # 响应PING消息
                pong_msg = self.protocol.create_pong(message.seq_id)
                await client.send_message(pong_msg)
                client.last_ping = time.time()

            elif message.msg_type == MessageType.DATA:
                # 回显数据消息
                echo_data = f"Echo: {message.data.decode('utf-8', errors='ignore')}"
                echo_msg = self.protocol.create_data(
                    seq_id=client.get_next_seq_id(),
                    data=echo_data.encode('utf-8')
                )
                await client.send_message(echo_msg)

                # 发送ACK确认
                ack_msg = self.protocol.create_ack(message.seq_id)
                await client.send_message(ack_msg)

            elif message.msg_type == MessageType.PONG:
                # 更新心跳时间
                client.last_ping = time.time()

            else:
                logger.warning(f"Unknown message type from {client.client_id}: {message.msg_type}")

        except Exception as e:
            logger.error(f"Error processing message from {client.client_id}: {e}")
            # 发送错误消息
            error_msg = self.protocol.create_error(
                seq_id=message.seq_id,
                error_msg=f"Processing error: {str(e)}"
            )
            try:
                await client.send_message(error_msg)
            except Exception as e:
                logger.error(f"send_message (error_msg) error -- {client.client_id}: {e}")

    async def cleanup_client(self, client_id):
        if client_id in self.clients:
            client = self.clients[client_id]
            await client.close()
            del self.clients[client_id]
            logger.info(f"Client {client_id} cleaned up")

    def setup_signal_handlers(self):
        """设置信号处理器"""

        def signal_handler(signum, frame):
            logger.info(f"Received signal {signum}")
            asyncio.create_task(self.stop_server())

        signal.signal(signal.SIGINT, signal_handler)
        signal.signal(signal.SIGTERM, signal_handler)

    async def get_server_stats(self) -> dict:
        """获取服务器统计信息"""
        return {
            'running': self.running,
            'client_count': len(self.clients),
            'clients': [
                {
                    'id': client.client_id,
                    'address': client.address,
                    'connected_at': client.connected_at,
                    'last_ping': client.last_ping
                }
                for client in self.clients.values()
            ]
        }

    async def heartbeat_monitor(self):
        """心跳监控任务"""
        while self.running:
            try:
                await asyncio.sleep(self.heartbeat_interval)

                current_time = time.time()
                timeout_clients = []

                # 检查超时的客户端
                for client_id, client in self.clients.items():
                    if current_time - client.last_ping > self.heartbeat_interval * 2:
                        timeout_clients.append(client_id)
                        logger.warning(f"Client {client_id} heartbeat timeout")

                # 清理超时客户端
                for client_id in timeout_clients:
                    await self.cleanup_client(client_id)

                # 向所有客户端发送心跳
                if self.clients:
                    ping_tasks = []
                    for client in self.clients.values():
                        ping_msg = self.protocol.create_ping(client.get_next_seq_id())
                        ping_tasks.append(client.send_message(ping_msg))

                    # 并发发送心跳，忽略失败的连接
                    await asyncio.gather(*ping_tasks, return_exceptions=True)

            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Heartbeat monitor error: {e}")

    async def stop_server(self):
        """停止服务器"""
        logger.info("Stopping server...")
        self.running = False

        # 停止心跳任务
        if self.heartbeat_task:
            self.heartbeat_task.cancel()
            try:
                await self.heartbeat_task
            except asyncio.CancelledError:
                pass

        # 关闭所有客户端连接
        if self.clients:
            logger.info(f"Closing {len(self.clients)} client connections...")
            close_tasks = [client.close() for client in self.clients.values()]
            await asyncio.gather(*close_tasks, return_exceptions=True)
            self.clients.clear()

        # 关闭服务器
        if self.server:
            self.server.close()
            await self.server.wait_closed()
            logger.info("Server stopped")
            self.server = None

    @asynccontextmanager
    async def run_context(self):
        """异步上下文管理器"""
        await self.start_server()
        try:
            yield self
        finally:
            await self.stop_server()


# 使用示例
async def main():
    """主函数"""
    server = AsyncTCPServer(host='0.0.0.0', port=8888)

    try:
        async with server.run_context():
            logger.info("Server is running. Press Ctrl+C to stop.")

            # 定期打印服务器状态
            while server.running:
                await asyncio.sleep(10)
                stats = await server.get_server_stats()
                logger.info(f"Server stats: {stats['client_count']} clients connected")

    except KeyboardInterrupt:
        logger.info("Server interrupted by user")
    except Exception as e:
        logger.error(f"Server error: {e}")


if __name__ == "__main__":
    asyncio.run(main())